
class ReservoirSampling:

    def __init__(self, vertex, para):
        self.budget = para['budget']
        self.vertex = vertex
        self.est = self.top_est = self.bottom_est = 1
        self.sampling = set()
        self.reverse = para['reverse']

    def add_vertex(self, vertex):
        changed = False
        '''
        for element in pair:
            if self.vertex != element.vertex:
                vertex = element
        '''
        if vertex.vertex not in self.sampling:
            if len(self.sampling) < self.budget:
                self.sampling.add(vertex)
                changed = True
            elif vertex.value < self.bottom_est:
                max_ele = self.get_max_sampling()
                self.sampling.remove(max_ele)
                if self.reverse:
                    max_ele.pop_sampled_by(self.vertex)
                self.sampling.add(vertex)
                self.top_est = max_ele.value
                self.bottom_est = vertex.value
                self.est = (self.top_est + self.bottom_est)/2
                changed = True
            if changed & self.reverse:
                vertex.sampled_by(self.vertex)
        return changed

    def get_max_sampling(self):
        value = 0
        max_element = None
        for vertex in self.sampling:
            if vertex.value > value:
                value = vertex.value
                max_element = vertex
        return max_element

    def get_count(self):
        return len(self.sampling)
